// $Id: wfs-loop.c 246 2013-10-21 05:07:48Z alan.watson.f $

#define _POSIX_C_SOURCE 200112L

#include <math.h>
#include <stdlib.h>

#include "oaxaca-config.h"
#include "oaxaca.h"

#include "array.h"
#include "detector.h"
#include "rfits.h"
#include "wfits.h"
#include "wfs.h"

#if 0

////////////////////////////////////////////////////////////////////////////////

static void
write_exposure(const char *fitsfilepath, array_short_t *a)
{
  oaxaca_info("writing FITS file.");

  FILE *fp = fopen(fitsfilepath, "wb");
  if (fp == NULL)
    oaxaca_error("unable to open \"%s\".", fitsfilepath);

  oaxaca_wfits_write_basic_header(fp, a->nx, a->ny);
  oaxaca_wfits_write_date(fp, "DATE");
  oaxaca_wfits_write_string(fp, "INSTRUME", "Oaxaca WFS");
  oaxaca_wfits_write_string(fp, "DETECTOR", "IDS uEye UI-5540SE-M S/N 4002849489");
  double dit  = 1.0 / oaxaca_shm->wfs.detector.framefrequency;
  oaxaca_wfits_write_double(fp, "DIT", dit, "%.6f");
  double ndit = oaxaca_shm->wfs.detector.cycle;
  oaxaca_wfits_write_double(fp, "NDIT", ndit, "%.0f");
  oaxaca_wfits_write_double(fp, "PIXFREQ", oaxaca_shm->wfs.detector.pixelfrequency, "%.1f");
  oaxaca_wfits_write_double(fp, "FRMFREQ", oaxaca_shm->wfs.detector.framefrequency, "%.3f");
  oaxaca_wfits_write_long(fp, "WINSX", oaxaca_shm->wfs.detector.window.sx);
  oaxaca_wfits_write_long(fp, "WINSY", oaxaca_shm->wfs.detector.window.sy);
  oaxaca_wfits_write_long(fp, "WINNX", oaxaca_shm->wfs.detector.window.nx);
  oaxaca_wfits_write_long(fp, "WINNY", oaxaca_shm->wfs.detector.window.ny);

  oaxaca_wfits_write_data(fp, a);

  fclose(fp);

  oaxaca_info("finished writing FITS file.");
}

////////////////////////////////////////////////////////////////////////////////

static array_long_t *sum = NULL;
static array_short_t *mean = NULL;

static void
coadd_start(void)
{
  oaxaca_info("exposing.");
    
  for (int iy = 0; iy < sum->ny; ++iy)
    for (int ix = 0; ix < sum->nx; ++ix)
      sum->z[iy][ix] = 0;    
}

static void
coadd_next(array_short_t *frame)
{
  for (int iy = 0; iy < frame->ny; ++iy)
    for (int ix = 0; ix < frame->nx; ++ix)
      sum->z[iy][ix] += frame->z[iy][ix];
}

static void
coadd_end()
{
  oaxaca_info("finished exposing.");

  long frames = oaxaca_shm->wfs.detector.cycle;
  for (int iy = 0; iy < mean->ny; ++iy)
    for (int ix = 0; ix < mean->nx; ++ix)
      mean->z[iy][ix] = sum->z[iy][ix] / frames;

  write_exposure("exposure.fits", mean);
}

////////////////////////////////////////////////////////////////////////////////

static array_short_t *dark = NULL;

static void
subtract_dark(array_short_t *frame)
{
  if (dark != NULL)
    for (int iy = 0; iy < frame->ny; ++iy)
      for (int ix = 0; ix < frame->nx; ++ix)
        frame->z[iy][ix] -= dark->z[iy][ix];
}

static void
write_dark(array_short_t *dark)
{
  const char *darkfitsfilepath = getdarkfitsfilepath();
  oaxaca_info("writing the dark image \"%s\".", darkfitsfilepath);
  write_exposure(darkfitsfilepath, dark);
  oaxaca_info("finished writing the dark image.");
}

static array_short_t *
read_dark(void)
{
  const char *darkfitsfilepath = getdarkfitsfilepath();

  oaxaca_info("reading the dark image \"%s\".", darkfitsfilepath);

  FILE *fp = fopen(darkfitsfilepath, "rb");
  if (fp == NULL) {
    oaxaca_info("unable to open the dark image.");
    return NULL;
  }
    
  array_short_t *dark = oaxaca_rfits(fp);
  fclose(fp);
  if (dark == NULL) {
    oaxaca_info("unable to read the dark image.");
    return NULL;
  }

  if (dark->nx != oaxaca_shm->wfs.detector.window.nx || 
      dark->ny != oaxaca_shm->wfs.detector.window.ny) {
    oaxaca_info("the dark image does not match the size of the wfs window.");
    free_array_short(dark);
    return NULL;
  }
  
  oaxaca_info("finished reading the dark image.");
  
  return dark;
}

////////////////////////////////////////////////////////////////////////////////

void
destripe(array_short_t *frame)
{
  long outersum = 0.0;
  for (int iy = 0; iy < frame->ny; ++iy) {
    long sum = 0;
    for (int ix = 0; ix < frame->nx / 4; ++ix)
      sum += frame->z[iy][ix];
    for (int ix = 3 * frame->nx / 4; ix < frame->nx; ++ix)
      sum += frame->z[iy][ix];
    long n = (frame->nx / 4) + (frame->nx - 3 * frame->nx / 4);
    short mean = sum / n;
    for (int ix = 0; ix < frame->nx; ++ix)
      frame->z[iy][ix] -= mean;
    outersum += mean;
  }
  short outermean = outersum / frame->ny;
  for (int iy = 0; iy < frame->ny; ++iy)
    for (int ix = 0; ix < frame->nx; ++ix)
      frame->z[iy][ix] += outermean;
}

////////////////////////////////////////////////////////////////////////////////

static double sum_dxc;
static double sum_dxc2;
static double sum_dyc;
static double sum_dyc2;
static double n;

static void
statistics_start(void)
{
  n = 0;
  sum_dxc = 0.0;
  sum_dxc2 = 0.0;
  sum_dyc = 0.0;
  sum_dyc2 = 0.0;
}

static void
statistics_next(double dxc, double dyc)
{
  n += 1.0;
  sum_dxc += dxc;
  sum_dxc2 += dxc * dxc;
  sum_dyc += dyc;
  sum_dyc2 += dyc * dyc;
}

static void
statistics_end(void)
{
  if (n != 0) {
    double mean_dxc = sum_dxc / n;
    double sigma_dxc = sqrt(sum_dxc2 / n - mean_dxc * mean_dxc);
    double mean_dyc = sum_dyc / n;
    double sigma_dyc = sqrt(sum_dyc2 / n - mean_dyc * mean_dyc);
    oaxaca_info("dxc: mean = %+6.2f sigma = %6.2f", mean_dxc, sigma_dxc);
    oaxaca_info("dyc: mean = %+6.2f sigma = %6.2f", mean_dyc, sigma_dyc);
  }
}

////////////////////////////////////////////////////////////////////////////////

static void
determine_spot_centroid(array_short_t *frame, oaxaca_spot_t *spot, double *yc, double *xc)
{
  size_t iyc0 = (int) spot->yc0;
  size_t ixc0 = (int) spot->xc0;
  size_t ms = spot->ms;
  size_t mc = spot->mc;

  size_t sys = iyc0 - ms;
  size_t sxs = ixc0 - ms;
  size_t eys = iyc0 + ms;
  size_t exs = ixc0 + ms;
  
  size_t iymax = sys;
  size_t ixmax = sxs;
  double zmax = frame->z[iymax][ixmax];
  for (size_t iy = sys; iy < eys; ++iy)
    for (size_t ix = sxs; ix < exs; ++ix)
      if (frame->z[iy][ix] > zmax) {
        iymax = iy;
        ixmax = ix;
        zmax = frame->z[iy][ix];
      }

  // Make sure the centroid region is within the spot region.
  
  size_t iyc;
  if (iymax < sys + mc)
    iyc = sys + mc;
  else if (iymax > eys - mc)
    iyc = eys - mc;
  else
    iyc = iymax;

  size_t ixc;
  if (ixmax < sxs + mc)
    ixc = sxs + mc;
  else if (ixmax > exs - mc)
    ixc = exs - mc;
  else
    ixc = ixmax;
  
  size_t syc = iyc - mc;
  size_t sxc = ixc - mc;
  size_t eyc = iyc + mc;
  size_t exc = ixc + mc;
  
  // Determine centroids in centroid region around spot.
  long w = 0, wx = 0, wy = 0;
  for (size_t iy = syc; iy <= eyc; ++iy)
    for (size_t ix = sxc; ix <= exc; ++ix) {
      w  += frame->z[iy][ix];
      wy += frame->z[iy][ix] * iy;
      wx += frame->z[iy][ix] * ix;
    }
    
  *yc = (double) wy / (double) w;  
  *xc = (double) wx / (double) w;
}

void
determine_error(array_short_t *frame, double *mean_dxc, double *mean_dyc)
{

  // Spots for (wfs-set-window! (- 732 144) 640 (* 2 288) 288)
  double x0 = 261.71 - 1.0 - 121.68;
  double y0 = 116.62 - 1.0 - 113.25;
  oaxaca_spot_t spot[] = {
    { .yc0 = y0 +  28.89, .xc0 = x0 + 119.39, .ms = 20, .mc = 10 },
    { .yc0 = y0 +  26.91, .xc0 = x0 + 169.52, .ms = 20, .mc = 10 },
    { .yc0 = y0 +  70.50, .xc0 = x0 +  73.55, .ms = 20, .mc = 10 },
    { .yc0 = y0 +  69.93, .xc0 = x0 + 119.93, .ms = 20, .mc = 10 },
    { .yc0 = y0 +  68.97, .xc0 = x0 + 168.06, .ms = 20, .mc = 10 },
    { .yc0 = y0 +  66.69, .xc0 = x0 + 213.51, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 113.45, .xc0 = x0 +  29.78, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 113.63, .xc0 = x0 +  75.31, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 113.25, .xc0 = x0 + 121.68, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 113.41, .xc0 = x0 + 168.79, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 112.52, .xc0 = x0 + 213.32, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 110.77, .xc0 = x0 + 256.87, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 158.32, .xc0 = x0 +  31.29, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 158.68, .xc0 = x0 +  76.42, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 158.68, .xc0 = x0 + 123.51, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 159.93, .xc0 = x0 + 170.66, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 160.08, .xc0 = x0 + 214.28, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 159.23, .xc0 = x0 + 256.58, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 201.21, .xc0 = x0 +  76.89, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 201.82, .xc0 = x0 + 128.40, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 206.91, .xc0 = x0 + 173.30, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 208.47, .xc0 = x0 + 216.29, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 253.46, .xc0 = x0 + 133.67, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 257.15, .xc0 = x0 + 178.52, .ms = 20, .mc = 10 },
  };
  size_t nspot = sizeof(spot) / sizeof(*spot);
  
  double dyc[nspot];
  double dxc[nspot];
  
  //oaxaca_info("%lu", (unsigned long) oaxaca_shm->wfs.detector.cycle);
  for (int i = 0; i < nspot; ++i) {
    double xc, yc;
    determine_spot_centroid(frame, &spot[i], &yc, &xc);
    dyc[i] = yc - spot[i].yc0;
    dxc[i] = xc - spot[i].xc0;
    //oaxaca_info("%2d %5.1f %5.1f %+5.1f %+5.1f ", i, yc, xc, dyc[i], dxc[i]);
  }
  
  *mean_dyc = 0.0;
  *mean_dxc = 0.0;
  for (int i = 0; i < nspot; ++i) {
    *mean_dyc += dyc[i];
    *mean_dxc += dxc[i];
  }
  *mean_dyc /= (double) nspot;
  *mean_dxc /= (double) nspot;

  //oaxaca_info("%lu %+6.2f %+6.2f", oaxaca_shm->wfs.detector.cycle, *mean_dxc, *mean_dyc);
}

////////////////////////////////////////////////////////////////////////////////

static void
correct(double mean_dxc, double mean_dyc)
{
  if (oaxaca_shm->wfc.command == 'I') {
    double aogain = oaxaca_shm->wfs.aogain;
    if (isfinite(mean_dxc))
      oaxaca_shm->wfc.mode[0] += aogain * mean_dxc;
    if (isfinite(mean_dyc))
      oaxaca_shm->wfc.mode[1] += aogain * mean_dyc;
    if (oaxaca_shm->wfs.detector.cycle % 40 == 0) {
      oaxaca_shm->wfc.mode[0] = 4.0 * ((double) rand() / (double) RAND_MAX - 0.5);
      rand(); rand(); rand();
      oaxaca_shm->wfc.mode[1] = 4.0 * ((double) rand() / (double) RAND_MAX - 0.5);
    }
    oaxaca_shm->wfc.modal = 1;
    oaxaca_shm->wfc.command = 't';
  }
}
  
////////////////////////////////////////////////////////////////////////////////

static unsigned long exposureframes;

int
loop_frame_handler(array_short_t *frame)
{
  if (oaxaca_shm->wfs.command == 'i') {
    oaxaca_info("aborting loop.");
    return 0;
  }

  if (oaxaca_shm->wfs.detector.cycle % exposureframes == 0) {
    statistics_start();
    coadd_start();
  }

  ++oaxaca_shm->wfs.detector.cycle;

  subtract_dark(frame);
  destripe(frame);
  coadd_next(frame);

  if (oaxaca_shm->wfs.command == 'A') {
    double mean_dxc, mean_dyc;
    determine_error(frame, &mean_dxc, &mean_dyc);
    correct(mean_dxc, mean_dyc);
    statistics_next(mean_dxc, mean_dyc);
  }
  
  if (oaxaca_shm->wfs.detector.cycle % exposureframes == 0) {
    statistics_end();
    coadd_end();
    if (oaxaca_shm->wfs.command != 'A')
      return 0;
  }

  return 1;
}

void
oaxaca_wfs_loop(void)
{
  oaxaca_info("looping.");

  oaxaca_shm->wfs.detector.cycle = 0;

  size_t nx = oaxaca_shm->wfs.detector.window.nx;
  size_t ny = oaxaca_shm->wfs.detector.window.ny;

  sum = alloc_array_long(ny, nx);
  mean = alloc_array_short(ny, nx);

  if (oaxaca_shm->wfs.command == 'D')
    dark = NULL;
  else
    dark = read_dark();

  coadd_start();
  if (oaxaca_shm->wfs.detector.exposuretime == 0.0)
    exposureframes = 1;
  else
    exposureframes = (unsigned long)
      ceil(oaxaca_shm->wfs.detector.exposuretime * oaxaca_shm->wfs.detector.framefrequency);

  oaxaca_ids_run(loop_frame_handler);

  if (oaxaca_shm->wfs.command == 'D')
    write_dark(mean);

  free_array_long(sum);
  free_array_short(mean); 
  if (dark != NULL)
    free_array_short(dark);

  oaxaca_info("finished looping.");
}

////////////////////////////////////////////////////////////////////////////////

#endif